{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EheA5_j_cEwc"
      },
      "source": [
        "##### Copyright 2019 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YCriMWd-pRTP"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OvRwFTkqcp1e"
      },
      "source": [
        "# Root Search\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/google/tf-quant-finance/blob/master/tf_quant_finance/examples/jupyter_notebooks/Root_Search.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4EL4whNjvUmV"
      },
      "outputs": [],
      "source": [
        "#@title Upgrade to TensorFlow 2.5+\n",
        "!pip install --upgrade tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 903,
          "status": "ok",
          "timestamp": 1629280526747,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "oF6t5bUovUmW",
        "outputId": "2e00b420-d0df-4cba-ecd4-b2dd7ce81ee7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Installing collected packages: tf-quant-finance\n",
            "Successfully installed tf-quant-finance-0.0.1.dev28\n"
          ]
        }
      ],
      "source": [
        "#@title Install TF Quant Finance\n",
        "!pip install tf-quant-finance"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lcQDncqBvUmW"
      },
      "source": [
        "### This notebook demonstrates the use of low level Tensorflow Quant Finance tools for root finding using Brent's method with emphasis on the following aspects:\n",
        "\n",
        "  * **Write Once**: Tensorflow supports GPUs without any significant code changes. The same model can be run on CPU and GPU\n",
        "  * **XLA Acceleration**: The **XLA compiler** can reduce overhead associated with the Tensorflow graph by fusing operators.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dr7lgj7mvUmW"
      },
      "outputs": [],
      "source": [
        "#@title Imports { display-mode: \"form\" }\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        " # tff for Tensorflow Finance\n",
        "import tf_quant_finance as tff \n",
        "\n",
        "root_search = tff.math.root_search\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\",\n",
        "                        category=FutureWarning)  # suppress printing warnings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 16,
          "status": "ok",
          "timestamp": 1629281741923,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "w3d7fZbDwivN",
        "outputId": "4218a7ea-0d0f-483b-c297-09a6c3f69b3b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Wed Aug 18 10:15:41 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 470.57.02    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   47C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lNSAPVK1vUmY"
      },
      "source": [
        "### Brent's Method\n",
        "Find the risk free rate for an interest rate swap:\n",
        "$$f(x) = \\log(\\sum_{i=0}^{N} e^{-r_i t_i}) - V_{swap}$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 6406,
          "status": "ok",
          "timestamp": 1629281315525,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "CHHfaVhNvUmY",
        "outputId": "9c03e31e-e2ab-4f75-eda3-b6f26623ec00"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "------------------------\n",
            "Tensorflow CPU (with auto-threading)\n",
            "Converged: tf.Tensor(True, shape=(), dtype=bool)\n",
            "Estimated root: tf.Tensor(0.4740769777681948, shape=(), dtype=float64)\n",
            "Objective at root: tf.Tensor(1.698336749011986e-08, shape=(), dtype=float64)\n",
            "Number of search steps: tf.Tensor(5, shape=(), dtype=int32)\n",
            "Timing:\n",
            "100 loops, best of 5: 3.36 ms per loop\n",
            "------------------------\n",
            "------------------------\n",
            "Tensorflow GPU\n",
            "Converged: tf.Tensor(True, shape=(), dtype=bool)\n",
            "Estimated root: tf.Tensor(0.4740769777681948, shape=(), dtype=float64)\n",
            "Objective at root: tf.Tensor(1.698336749011986e-08, shape=(), dtype=float64)\n",
            "Number of search steps: tf.Tensor(5, shape=(), dtype=int32)\n",
            "Timing:\n",
            "100 loops, best of 5: 6.66 ms per loop\n",
            "------------------------\n"
          ]
        }
      ],
      "source": [
        "#@title Search Range\n",
        "\n",
        "number_of_tenors = 100  #@param\n",
        "swap_value = 0.5  #@param\n",
        "\n",
        "dtype = tf.float64\n",
        "\n",
        "tenors = range(1, number_of_tenors + 1)\n",
        "\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "# Wrap with TF function for better performance\n",
        "root_search_tf = tf.function(root_search.brentq)\n",
        "\n",
        "def run_on_device(device):\n",
        "  with tf.device(device):\n",
        "    return root_search_tf(\n",
        "      valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))\n",
        "\n",
        "## TFF on CPU\n",
        "brent_result = run_on_device('/cpu:0')\n",
        "\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = brent_result\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow CPU (with auto-threading)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit -n 100 run_on_device('/cpu:0')\n",
        "print(\"------------------------\")\n",
        "\n",
        "## TFF on GPU\n",
        "brent_result = run_on_device('/gpu:0')\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = brent_result\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow GPU\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit -n 100 run_on_device('/gpu:0')\n",
        "print(\"------------------------\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U_CHQQN2vUmY"
      },
      "source": [
        "### Speedup from XLA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 1687,
          "status": "ok",
          "timestamp": 1629281355775,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "SSvHg-LivUmZ",
        "outputId": "567785c5-06de-4e9a-8bb3-d710012e5828"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "------------------------\n",
            "Tensorflow CPU (compiled with XLA)\n",
            "Converged: tf.Tensor(True, shape=(), dtype=bool)\n",
            "Estimated root: tf.Tensor(0.47407697776819474, shape=(), dtype=float64)\n",
            "Objective at root: tf.Tensor(1.6983367823186768e-08, shape=(), dtype=float64)\n",
            "Number of search steps: tf.Tensor(5, shape=(), dtype=int32)\n",
            "Timing:\n",
            "100 loops, best of 5: 458 µs per loop\n",
            "------------------------\n",
            "------------------------\n",
            "Tensorflow GPU (compiled with XLA)\n",
            "Converged: tf.Tensor(True, shape=(), dtype=bool)\n",
            "Estimated root: tf.Tensor(0.4740769777681948, shape=(), dtype=float64)\n",
            "Objective at root: tf.Tensor(1.6983367712164466e-08, shape=(), dtype=float64)\n",
            "Number of search steps: tf.Tensor(5, shape=(), dtype=int32)\n",
            "Timing:\n",
            "100 loops, best of 5: 834 µs per loop\n",
            "------------------------\n"
          ]
        }
      ],
      "source": [
        "#@title Search Range\n",
        "\n",
        "number_of_tenors = 100  #@param\n",
        "swap_value = 0.5  #@param\n",
        "\n",
        "dtype = tf.float64\n",
        "\n",
        "tenors = range(1, number_of_tenors + 1)\n",
        "\n",
        "tf.compat.v1.reset_default_graph()\n",
        "\n",
        "t = tf.constant(tenors, dtype=dtype)\n",
        "v = tf.constant(swap_value, dtype=dtype)\n",
        "\n",
        "def valuation_fn(x):\n",
        "  return tf.reduce_logsumexp(-x * t) - v\n",
        "\n",
        "# Wrap with TF function for better performance\n",
        "root_search_xla = tf.function(root_search.brentq, jit_compile=True)\n",
        "\n",
        "def run_on_device(device):\n",
        "  with tf.device(device):\n",
        "    return root_search_xla(\n",
        "      valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))\n",
        "\n",
        "## TFF on CPU compiled with XLA\n",
        "brent_result = run_on_device('/cpu:0')\n",
        "\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = brent_result\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow CPU (compiled with XLA)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit -n 100 brent_result = run_on_device('/cpu:0')\n",
        "\n",
        "print(\"------------------------\")\n",
        "\n",
        "## TFF on GPU compiled with XLA\n",
        "brent_result = run_on_device('/gpu:0')\n",
        "\n",
        "\n",
        "estimated_root, objective_at_root, num_iterations, converged = brent_result\n",
        "\n",
        "print(\"------------------------\")\n",
        "print(\"Tensorflow GPU (compiled with XLA)\")\n",
        "print(\"Converged:\", converged)\n",
        "print(\"Estimated root:\", estimated_root)\n",
        "print(\"Objective at root:\", objective_at_root)\n",
        "print(\"Number of search steps:\", num_iterations)\n",
        "print(\"Timing:\")\n",
        "%timeit -n 100 brent_result = run_on_device('/gpu:0')\n",
        "print(\"------------------------\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "private"
      },
      "name": "Root_Search.ipynb",
      "provenance": [
        {
          "file_id": "1rT50234_4wb5VFNWyheBfyVlV-iuRA0t",
          "timestamp": 1553699488325
        },
        {
          "file_id": "1G2Wm6NzpdH3gPFsFTaLslgonTTkI0U0h",
          "timestamp": 1552493606270
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
